package tree.迭代bfs;

import po.TreeNode;

import java.util.*;

/**
 * @author songZiHao
 * @version 1.0.0
 * @ClassName N二叉树中所有距离为K的结点863.java
 * @Description https://leetcode-cn.com/problems/all-nodes-distance-k-in-binary-tree/
 * @createTime 2021年07月20日 19:34:00
 */
public class N二叉树中所有距离为K的结点863 {
    /**
     * 这里的思路是，先记录每个节点的父节点信息，在后面进行从下往上的层序遍历的时候会用到
     *
     * 1. 对目标节点进行层序遍历
     * 2. 每个遍历到的节点需要加入到对应的已经遍历节点集合中
     * 3. 进行遍历的时候需要判断对应的节点是不是已经遍历到了，避免重复遍历
     * 4. 进行从上往下层序遍历的同时，也需要进行从下往上的层序遍历
     *
     */

    private Map<TreeNode, TreeNode> parents;
    private Set<TreeNode> visited;
    private List<Integer> ans;
    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        parents = new HashMap<>();
        //获取每个节点对应的父节点信息，用于后面的向上遍历
        getParent(root, null);
        visited = new HashSet<>();
        ans = new ArrayList<>();
        Deque<TreeNode> queue = new LinkedList<>();
        queue.add(target);
        //广度优先，从目标节点开始找距离为k的节点
        while (!queue.isEmpty()) {
            //层序遍历一层就减去一层
            if (K-- == 0) {
                while (!queue.isEmpty()) {
                    ans.add(queue.pop().val);
                }
                continue;
            }
            int size = queue.size();
            for (int i = 0; i < size; i++) {
                TreeNode node = queue.poll();
                //节点加入已经遍历过的集合中
                visited.add(node);
                //如果左节点不为空，并且节点没有遍历过，就加入到层序遍历的队列中。
                //这里判断有没有遍历过的原因是因为，同时进行向上遍历的时候，会无法避免的对父节点进行向下的层序遍历，如果加入到集合，就可以知道这一层是不是已经遍历过了
                if (node.left != null && !visited.contains(node.left)) {
                    queue.add(node.left);
                }
                //处理右节点
                if (node.right != null && !visited.contains(node.right)) {
                    queue.add(node.right);
                }
                //从整个树中拿到对应节点的父节点，这就从目标节点层向上遍历的逻辑
                TreeNode p = parents.get(node);
                //如果父节点不为空，则加入到队列中，
                if (p != null && !visited.contains(p)) {
                    queue.add(p);
                }
            }
        }
        return ans;
    }

    public void getParent(TreeNode node, TreeNode par) {
        if (node != null) {
            parents.put(node, par);
            getParent(node.left, node);
            getParent(node.right, node);
        }
    }
}
